Hi @sarthakpati and team,

Thanks for this nice work! I wrote a quick PoC to showcase that you can easily have integration with the 🤗 hub so that you can automatically load the various GaNDLF models using from_pretrained (and push them using push_to_hub), track download numbers for your models (similar to models in the Transformers library), and have nice model cards on a per-model basis. It leverages the PyTorchModelHubMixin class which allows to inherits these methods.

Usage is as follows:

from .models import DenseNet

model = DenseNet(...)

# optionally, one can push a trained model to the hub at the end of training
model.push_to_hub("nielsr/densenet-finetuned-medical-segmentation")

# reload as follows
model = DenseNet.from_pretrained("nielsr/densenet-finetuned-medical-segmentation")

Alternatively, if you want a more custom integration with the hub, we also provide methods like hf_hub_download (to download a file from the hub, upload_folder and upload_file). Those require some more work however to make download stats work.

This means people don't need to manually download a checkpoint first in their local environment, it just loads automatically from the hub.

Would you be interested in this integration?

Fixes #727

Kind regards,

Niels